""""""
import time

from datetime import datetime
from pytz import timezone
from typing import List

from pymongo import ASCENDING, MongoClient, ReplaceOne
from pymongo.database import Database
from pymongo.cursor import Cursor
from pymongo.collection import Collection
from pymongo.results import DeleteResult

from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import BarData, TickData, FutBasicData, FutMappingData
from vnpy.trader.database import BaseDatabase, BarOverview, DB_TZ
from vnpy.trader.setting import SETTINGS

from pandas import DataFrame
import json


class MongodbDatabase(BaseDatabase):
    """MongoDB数据库接口"""

    def __init__(self) -> None:
        """"""
        # 读取配置
        self.database: str = SETTINGS["database.database"]
        self.host: str = SETTINGS["database.host"]
        self.port: int = SETTINGS["database.port"]
        self.username: str = SETTINGS["database.user"]
        self.password: str = SETTINGS["database.password"]

        # 创建客户端
        if self.username and self.password:
            self.client: MongoClient = MongoClient(
                host=self.host,
                port=self.port,
                tz_aware=True,
                username=self.username,
                password=self.password,
                tzinfo=DB_TZ
            )
        else:
            self.client: MongoClient = MongoClient(
                host=self.host,
                port=self.port,
                tz_aware=True,
                tzinfo=DB_TZ
            )

        # 初始化数据库
        self.db: Database = self.client[self.database]

        # 初始化K线数据表
        self.bar_collection: Collection = self.db["bar_data"]
        self.bar_collection.create_index(
            [
                ("exchange", ASCENDING),
                ("symbol", ASCENDING),
                ("interval", ASCENDING),
                ("datetime", ASCENDING),
            ],
            unique=True
        )

        # 初始化Tick数据表
        self.tick_collection: Collection = self.db["tick_data"]
        self.tick_collection.create_index(
            [
                ("exchange", ASCENDING),
                ("symbol", ASCENDING),
                ("datetime", ASCENDING),
            ],
            unique=True
        )

        # 初始化K线概览表
        self.overview_collection: Collection = self.db["bar_overview"]
        self.overview_collection.create_index(
            [
                ("exchange", ASCENDING),
                ("symbol", ASCENDING),
                ("interval", ASCENDING),
            ],
            unique=True
        )

        # 初始化fut_basic数据表
        self.fut_basic_collection: Collection = self.db["fut_basic"]
        self.fut_basic_collection.create_index(
            [
                ("exchange", ASCENDING),
                ("symbol", ASCENDING),
            ],
            unique=True
        )

        # 初始化fut_mapping数据表
        self.fut_mapping_collection: Collection = self.db["fut_mapping"]
        self.fut_mapping_collection.create_index(
            [
                ("ts_code", ASCENDING),
                ("trade_date", ASCENDING),
            ],
            unique=True
        )

    def save_bar_data_pd(self, bars: DataFrame) -> bool:
        """保存K线数据"""
        requests: List[ReplaceOne] = []

        for bar in bars.itertuples():
            # 逐个插入
            filter: dict = {
                "symbol": bar.symbol,
                "exchange": bar.exchange.value,
                "datetime": bar.datetime,
                "interval": bar.interval.value,
            }

            d: dict = {
                "symbol": bar.symbol,
                "exchange": bar.exchange.value,
                "datetime": bar.datetime,
                "interval": bar.interval.value,
                "volume": bar.volume,
                "turnover": bar.turnover,
                "open_interest": bar.open_interest,
                "open_price": bar.open,
                "high_price": bar.high,
                "low_price": bar.low,
                "close_price": bar.close,
            }
            requests.append(ReplaceOne(filter, d, upsert=True))

        # bar_data = json.loads(bars.T.to_json()).values()
        # result = self.bar_collection.insert_many(bar_data)

        self.bar_collection.bulk_write(requests, ordered=False)

        # 更新汇总
        filter: dict = {
            "symbol": bar.symbol,
            "exchange": bar.exchange.value,
            "interval": bar.interval.value
        }

        overview: dict = self.overview_collection.find_one(filter)
        print(bars['datetime'].min())
        # datetime.timezone()
        # print(overview["start"])

        if not overview:
            overview = {
                "symbol": bar.symbol,
                "exchange": bar.exchange.value,
                "interval": bar.interval.value,
                "count": len(bars),
                "start": bars['datetime'].min(),
                "end": bars['datetime'].max()
            }
        else:
            # timezone('Asia/Shanghai')
            overview["start"] = min(bars['datetime'].min(), overview["start"].replace(tzinfo=None))
            overview["end"] = max(bars['datetime'].max(), overview["end"].replace(tzinfo=None))
            overview["count"] = self.bar_collection.count_documents(filter)

        self.overview_collection.update_one(filter, {"$set": overview}, upsert=True)

        return True



    def save_tick_data(self, ticks: List[TickData]) -> bool:
        """保存TICK数据"""
        requests: List[ReplaceOne] = []

        for tick in ticks:
            filter: dict = {
                "symbol": tick.symbol,
                "exchange": tick.exchange.value,
                "datetime": tick.datetime,
            }

            d: dict = {
                "symbol": tick.symbol,
                "exchange": tick.exchange.value,
                "datetime": tick.datetime,
                "name": tick.name,
                "volume": tick.volume,
                "turnover": tick.turnover,
                "open_interest": tick.open_interest,
                "last_price": tick.last_price,
                "last_volume": tick.last_volume,
                "limit_up": tick.limit_up,
                "limit_down": tick.limit_down,
                "open_price": tick.open_price,
                "high_price": tick.high_price,
                "low_price": tick.low_price,
                "pre_close": tick.pre_close,
                "bid_price_1": tick.bid_price_1,
                "bid_price_2": tick.bid_price_2,
                "bid_price_3": tick.bid_price_3,
                "bid_price_4": tick.bid_price_4,
                "bid_price_5": tick.bid_price_5,
                "ask_price_1": tick.ask_price_1,
                "ask_price_2": tick.ask_price_2,
                "ask_price_3": tick.ask_price_3,
                "ask_price_4": tick.ask_price_4,
                "ask_price_5": tick.ask_price_5,
                "bid_volume_1": tick.bid_volume_1,
                "bid_volume_2": tick.bid_volume_2,
                "bid_volume_3": tick.bid_volume_3,
                "bid_volume_4": tick.bid_volume_4,
                "bid_volume_5": tick.bid_volume_5,
                "ask_volume_1": tick.ask_volume_1,
                "ask_volume_2": tick.ask_volume_2,
                "ask_volume_3": tick.ask_volume_3,
                "ask_volume_4": tick.ask_volume_4,
                "ask_volume_5": tick.ask_volume_5,
                "localtime": tick.localtime,
            }

            requests.append(ReplaceOne(filter, d, upsert=True))

        self.tick_collection.bulk_write(requests, ordered=False)

        return True

    def save_fut_basic(self, futs: List[FutBasicData]) -> bool:
        """保存K线数据"""
        requests: List[ReplaceOne] = []

        for fut in futs:
            # 逐个插入
            filter: dict = {
                "symbol": fut.symbol,
                "exchange": fut.exchange.value,
            }

            d: dict = {
                "symbol": fut.symbol, # 合约代码
                "exchange": fut.exchange.value, # 交易市场
                "ts_code": fut.ts_code,  # 合约代码
                "name": fut.name,  # 中文简称
                "fut_code": fut.fut_code, # 合约产品代码
                "multiplier": fut.multiplier,  # 合约乘数
                "trade_unit": fut.trade_unit, # 交易计量单位
                "per_unit": fut.per_unit, # 交易单位(每手)
                "quote_unit": fut.quote_unit,  # 报价单位
                "quote_unit_desc": fut.quote_unit_desc, # 最小报价单位说明
                "d_mode_desc" : fut.d_mode_desc,  # 交割方式说明
                "list_date" : fut.list_date,  # 上市日期
                "delist_date" : fut.delist_date,  # 最后交易日期
                "d_month" : fut.d_month,  # 交割月份
                "last_ddate" : fut.last_ddate,  # 最后交割日
                "trade_time_desc" : fut.trade_time_desc,  # 交易时间说明
            }

            requests.append(ReplaceOne(filter, d, upsert=True))

        self.fut_basic_collection.bulk_write(requests, ordered=False)

        return True

    def save_fut_mapping(self, futs: List[FutMappingData]) -> bool:
        """保存K线数据"""
        requests: List[ReplaceOne] = []

        for fut in futs:
            # 逐个插入
            filter: dict = {
                "ts_code": fut.ts_code,
                "trade_date": fut.trade_date,
            }

            d: dict = {
                "ts_code": fut.ts_code,  # 合约代码
                "trade_date": fut.trade_date,  # 交易日期
                "mapping_ts_code": fut.mapping_ts_code,  # 期货合约代码
            }

            requests.append(ReplaceOne(filter, d, upsert=True))

        self.fut_mapping_collection.bulk_write(requests, ordered=False)

        return True

    def load_bar_data(
        self,
        symbol: str,
        exchange: Exchange,
        interval: Interval,
        start: datetime,
        end: datetime
    ) -> List[BarData]:
        """读取K线数据"""
        filter: dict = {
            "symbol": symbol,
            "exchange": exchange.value,
            "interval": interval.value,
            "datetime": {
                "$gte": start.astimezone(DB_TZ),
                "$lte": end.astimezone(DB_TZ)
            }
        }

        c: Cursor = self.bar_collection.find(filter)

        bars: List[BarData] = []
        for d in c:
            d["exchange"] = Exchange(d["exchange"])
            d["interval"] = Interval(d["interval"])
            d["gateway_name"] = "DB"
            d.pop("_id")

            bar = BarData(**d)
            bars.append(bar)

        return bars

    def load_tick_data(
        self,
        symbol: str,
        exchange: Exchange,
        start: datetime,
        end: datetime
    ) -> List[TickData]:
        """读取TICK数据"""
        filter: dict = {
            "symbol": symbol,
            "exchange": exchange.value,
            "datetime": {
                "$gte": start.astimezone(DB_TZ),
                "$lte": end.astimezone(DB_TZ)
            }
        }

        c: Cursor = self.tick_collection.find(filter)

        ticks: List[TickData] = []
        for d in c:
            d["exchange"] = Exchange(d["exchange"])
            d["gateway_name"] = "DB"
            d.pop("_id")

            tick: TickData = TickData(**d)
            ticks.append(tick)

        return ticks

    def delete_bar_data(
        self,
        symbol: str,
        exchange: Exchange,
        interval: Interval
    ) -> int:
        """删除K线数据"""
        filter: dict = {
            "symbol": symbol,
            "exchange": exchange.value,
            "interval": interval.value,
        }

        result: DeleteResult = self.bar_collection.delete_many(filter)
        self.overview_collection.delete_one(filter)

        return result.deleted_count

    def delete_tick_data(
        self,
        symbol: str,
        exchange: Exchange
    ) -> int:
        """删除TICK数据"""
        filter: dict = {
            "symbol": symbol,
            "exchange": exchange.value
        }

        result: DeleteResult = self.tick_collection.delete_many(filter)
        return result.deleted_count

    def get_bar_overview(self) -> List[BarOverview]:
        """查询数据库中的K线汇总信息"""
        c: Cursor = self.overview_collection.find()

        overviews: List[BarOverview] = []
        for d in c:
            d["exchange"] = Exchange(d["exchange"])
            d["interval"] = Interval(d["interval"])
            d.pop("_id")

            overview: BarOverview = BarOverview(**d)
            overviews.append(overview)

        return overviews

    def quary_fut_basic(self, fut_code: str = "SA", fut_type: str = "2") -> List[FutBasicData]:
        """
        Query fut basic from datafeed.
        """
        # fut_type 0，所有 1:主力 2：连续 3：普通
        filter: dict = {
            "fut_code": fut_code,
            # "name": {'$regex': '主力'}
        }
        if(fut_type == "1"):
            filter: dict = {
                "fut_code": fut_code,
                "name": {'$regex': '主力'}
            }
        elif (fut_type == "2"):
            filter: dict = {
                "fut_code": fut_code,
                "name": {'$regex': '连续'}
            }
        else:
            filter: dict = {
                "fut_code": fut_code,
            }

        c: Cursor = self.fut_basic_collection.find(filter)

        datas: List[FutBasicData] = []
        for d in c:
            d["exchange"] = Exchange(d["exchange"])
            # d["interval"] = Interval(d["interval"])
            d["gateway_name"] = "DB"
            d.pop("_id")

            data = FutBasicData(**d)
            datas.append(data)

        return datas
